#############################################################################################
#############################################################################################
# 
#   Open Source License/Disclaimer, Forecast Systems Laboratory NOAA/OAR/GSD, 
#   325 Broadway Boulder, CO 80305
#
#   This software is distributed under the Open Source Definition, which may be 
#   found at http://www.opensource.org/.
#   In particular, redistribution and use in source and binary forms, with or 
#   without modification, are permitted provided that the following conditions are met:
#
#   - Redistributions of source code must retain this notice, this list of 
#     conditions and the following disclaimer.
#
#   - Redistributions in binary form must provide access to this notice, this 
#     list of conditions and the following disclaimer, and the underlying source code.
#
#   - All modifications to this software must be clearly documented, and are 
#     solely the responsibility of the agent making the modifications.
#
#   - If significant modifications or enhancements are made to this software, 
#     the GSD Software Policy Manager (softwaremgr.fsl@noaa.gov) should be notified.
#
#   THIS SOFTWARE AND ITS DOCUMENTATION ARE IN THE PUBLIC DOMAIN AND ARE 
#   FURNISHED "AS IS." THE AUTHORS, THE UNITED STATES GOVERNMENT, ITS INSTRUMENTALITIES,
#   OFFICERS, EMPLOYEES, AND AGENTS MAKE NO WARRANTY, EXPRESS OR IMPLIED, AS TO
#   THE USEFULNESS OF THE SOFTWARE AND DOCUMENTATION FOR ANY PURPOSE. THEY ASSUME 
#   NO RESPONSIBILITY (1) FOR THE USE OF THE SOFTWARE AND DOCUMENTATION; OR (2) TO PROVIDE
#   TECHNICAL SUPPORT TO USERS. 
#############################################################################################
#############################################################################################


import os
import numpy as np
from cdar.models.Shallow_Water_2D.ShallowWater2DModel import ShallowWaterModelClass as model
import cdar.models.Shallow_Water_2D.ShallowWater2DModelVis as vis
from cdar.utilities import namelist_read as nml
  
DataFilePath='./data'  # Path to store or read FG, B and OBS

# Read in parameters for model configurations
params = nml.namelist_read(DataFilePath)

domain = model(params)  

# Read Truth.
print "Reading TRUTH from file %s" %  os.path.join(DataFilePath,'TRUTH_'+domain.Name+'_'+params['start_time']+'.npz')
npzfile = np.load(os.path.join(DataFilePath,'TRUTH_'+domain.Name+'_'+params['start_time']+'.npz'))
domain.ic = npzfile['TRUTH']
#u,v,h = np.split(domain.ic, domain.VarDim)
#vis.plot_3d(domain.x,domain.y,h,'Truth')

# Propogate the model
truth_traj = domain.model_integration()

# Read FG.
print "Reading FG from file %s" %  os.path.join(DataFilePath,'FG_'+domain.Name+'_'+params['start_time']+'.npz')
npzfile = np.load(os.path.join(DataFilePath,'FG_'+domain.Name+'_'+params['start_time']+'.npz'))
domain.ic = npzfile['FG']
#fg_u,fg_v,fg_h = np.split(domain.ic, domain.VarDim)
#vis.plot_3d(domain.x,domain.y,fg_h,'FG: rmse (u,v,h) = %6.3e %6.3e %6.3e' % (np.sqrt(np.sum((fg_u - u) ** 2)/u.size),\
#           np.sqrt(np.sum((fg_v - v) ** 2)/v.size),np.sqrt(np.sum((fg_h - h) ** 2)/h.size)) )

# Propogate the model
fg_traj = domain.model_integration()
  
# Read analysis
print "Reading Analysis from file %s" %  os.path.join(DataFilePath,'Analysis_'+domain.Name+'_'+params['start_time']+'.npz')
npzfile = np.load(os.path.join(DataFilePath,'Analysis_'+domain.Name+'_'+params['start_time']+'.npz'))
domain.ic = npzfile['Analysis']
#ana_u,ana_v,ana_h = np.split(domain.ic, domain.VarDim)
#vis.plot_3d(domain.x,domain.y,ana_h,'Analysis: rmse (u,v,h) = %6.3e %6.3e %6.3e' % (np.sqrt(np.sum((ana_u - u) ** 2)/u.size),\
#           np.sqrt(np.sum((ana_v - v) ** 2)/v.size),np.sqrt(np.sum((ana_h - h) ** 2/h.size))) )

# Propogate the model
ana_traj = domain.model_integration()

# Read Observations
obs_list = []
obs_truth_list = []
innov_time = []

for time in domain.history:  
    if os.path.exists(os.path.join(DataFilePath,'OBS_'+domain.Name+'_'+str(int(time))+'.npz')) :
        print 'Reading observation file at %s from %s' % ( str(time), os.path.join(DataFilePath,'OBS_'+domain.Name+'_'+str(int(time))+'.npz'))
        npzfile = np.load(os.path.join(DataFilePath,'OBS_'+domain.Name+'_'+str(int(time))+'.npz'))
        obs_list.append(npzfile['OBS'])
        obs_error = npzfile['OBS_ERROR'] 
        
        print 'Reading TRUTH observation file at %s from %s' % ( str(time), os.path.join(DataFilePath,'OBS_TRUTH_'+domain.Name+'_'+str(int(time))+'.npz'))
        npzfile = np.load(os.path.join(DataFilePath,'OBS_TRUTH_'+domain.Name+'_'+str(int(time))+'.npz'))
        obs_truth_list.append(npzfile['OBS'])

        # Record the time when the observation is available
        innov_time.append(time)

# Calculate RMSE of Truth, FG and Analysis on observations.
# Calculate the innovations 
for time in domain.history:
        
    if time in innov_time:
                
        print "=====================Time = %d======================" % time
        
        ###################
        #  For Truth    
        ###################
           
        # Ploting
        u,v,h = np.split(truth_traj[np.where(domain.history == time),:].reshape(domain.ic.shape),domain.VarDim)
        vis.plot_3d(domain.x, domain.y, h, 'True H at Time ='+str(time))
        
        # Strip off the boundaries of traj for DA
        hx = domain.strip_off_boundaries(truth_traj[np.where(domain.history == time),:].reshape(domain.ic.shape))
    
        # H(X) - Y
        innovation = hx - obs_truth_list[innov_time.index(time)]
        u_innov, v_innov, h_innov =np.split(innovation, domain.VarDim)
        u_obs_number = u_innov[~np.isnan(u_innov)].size
        v_obs_number = v_innov[~np.isnan(v_innov)].size
        h_obs_number = h_innov[~np.isnan(h_innov)].size
        u_innov[np.isnan(u_innov)] = 0.0
        v_innov[np.isnan(v_innov)] = 0.0
        h_innov[np.isnan(h_innov)] = 0.0
        
        # Calculate RMSE
        u_RMSE = np.sqrt( np.dot( u_innov.reshape(u_innov.size).transpose(), u_innov.reshape(u_innov.size) )/u_obs_number )
        v_RMSE = np.sqrt( np.dot( v_innov.reshape(v_innov.size).transpose(), v_innov.reshape(v_innov.size) )/v_obs_number )
        h_RMSE = np.sqrt( np.dot( h_innov.reshape(h_innov.size).transpose(), h_innov.reshape(h_innov.size) )/h_obs_number )
              
        print "TRUTH  Observation         (%d, %d, %d) verification RMSE(u,v,h) : %f, %f, %f " % ( u_obs_number, v_obs_number, h_obs_number, u_RMSE,v_RMSE, h_RMSE)
         
        ###################
        #  For Observation       
        ###################
        
        # Obs- True_Obs
        innovation = obs_list[innov_time.index(time)] - obs_truth_list[innov_time.index(time)]
        u_innov, v_innov, h_innov =np.split(innovation, domain.VarDim)
        u_obs_number = u_innov[~np.isnan(u_innov)].size
        v_obs_number = v_innov[~np.isnan(v_innov)].size
        h_obs_number = h_innov[~np.isnan(h_innov)].size
        u_innov[np.isnan(u_innov)] = 0.0
        v_innov[np.isnan(v_innov)] = 0.0
        h_innov[np.isnan(h_innov)] = 0.0
        
        # Calculate RMSE
        u_RMSE = np.sqrt( np.dot( u_innov.reshape(u_innov.size).transpose(), u_innov.reshape(u_innov.size) )/u_obs_number )
        v_RMSE = np.sqrt( np.dot( v_innov.reshape(v_innov.size).transpose(), v_innov.reshape(v_innov.size) )/v_obs_number )
        h_RMSE = np.sqrt( np.dot( h_innov.reshape(h_innov.size).transpose(), h_innov.reshape(h_innov.size) )/h_obs_number )
              
        print "Observation error          (%d, %d, %d) verification RMSE(u,v,h) : %f, %f, %f " % ( u_obs_number, v_obs_number, h_obs_number, u_RMSE,v_RMSE, h_RMSE)
        
        ###################
        #  For First Guess
        ###################

        # Ploting
        fg_u,fg_v,fg_h = np.split(fg_traj[np.where(domain.history == time),:].reshape(domain.ic.shape),domain.VarDim)
        vis.plot_3d(domain.x,domain.y,fg_h,'FG H at Time ='+str(time)+"\n"+'rmse (u,v,h) = %6.3e, %6.3e, %6.3e' % (np.sqrt(np.sum((fg_u - u) ** 2)/u.size),\
           np.sqrt(np.sum((fg_v - v) ** 2)/v.size),np.sqrt(np.sum((fg_h - h) ** 2)/h.size)) )
        
        # Strip off the boundaries of traj for DA
        hx = domain.strip_off_boundaries(fg_traj[np.where(domain.history == time),:].reshape(domain.ic.shape))
    
        # H(X) - Y
        innovation = hx - obs_truth_list[innov_time.index(time)]
        u_innov, v_innov, h_innov =np.split(innovation, domain.VarDim)
        u_obs_number = u_innov[~np.isnan(u_innov)].size
        v_obs_number = v_innov[~np.isnan(v_innov)].size
        h_obs_number = h_innov[~np.isnan(h_innov)].size
        u_innov[np.isnan(u_innov)] = 0.0
        v_innov[np.isnan(v_innov)] = 0.0
        h_innov[np.isnan(h_innov)] = 0.0
        
        # Calculate RMSE
        u_RMSE = np.sqrt( np.dot( u_innov.reshape(u_innov.size).transpose(), u_innov.reshape(u_innov.size) )/u_obs_number )
        v_RMSE = np.sqrt( np.dot( v_innov.reshape(v_innov.size).transpose(), v_innov.reshape(v_innov.size) )/v_obs_number )
        h_RMSE = np.sqrt( np.dot( h_innov.reshape(h_innov.size).transpose(), h_innov.reshape(h_innov.size) )/h_obs_number )
                            
        print "FG     against Observation (%d, %d, %d) verification RMSE(u,v,h) : %f, %f, %f " % ( u_obs_number, v_obs_number, h_obs_number, u_RMSE,v_RMSE, h_RMSE)
        
        ###################
        #  For Analysis 
        ###################
        
        # Ploting
        ana_u,ana_v,ana_h = np.split(ana_traj[np.where(domain.history == time),:].reshape(domain.ic.shape),domain.VarDim)
        vis.plot_3d(domain.x,domain.y,ana_h,'Analysis at Time ='+str(time)+"\n"+'rmse (u,v,h) = %6.3e, %6.3e, %6.3e' % (np.sqrt(np.sum((ana_u - u) ** 2)/u.size),\
           np.sqrt(np.sum((ana_v - v) ** 2)/v.size),np.sqrt(np.sum((ana_h - h) ** 2/h.size))) )

        # Strip off the boundaries of traj for DA
        hx = domain.strip_off_boundaries(ana_traj[np.where(domain.history == time),:].reshape(domain.ic.shape))
    
        # H(X) - Y
        innovation = hx - obs_truth_list[innov_time.index(time)]
        u_innov, v_innov, h_innov =np.split(innovation, domain.VarDim)
        u_obs_number = u_innov[~np.isnan(u_innov)].size
        v_obs_number = v_innov[~np.isnan(v_innov)].size
        h_obs_number = h_innov[~np.isnan(h_innov)].size
        u_innov[np.isnan(u_innov)] = 0.0
        v_innov[np.isnan(v_innov)] = 0.0
        h_innov[np.isnan(h_innov)] = 0.0
            
        # Calculate RMSE
        u_RMSE = np.sqrt( np.dot( u_innov.reshape(u_innov.size).transpose(), u_innov.reshape(u_innov.size) )/u_obs_number )
        v_RMSE = np.sqrt( np.dot( v_innov.reshape(v_innov.size).transpose(), v_innov.reshape(v_innov.size) )/v_obs_number )
        h_RMSE = np.sqrt( np.dot( h_innov.reshape(h_innov.size).transpose(), h_innov.reshape(h_innov.size) )/h_obs_number )
                     
                
        print "Analy. against Observation (%d, %d, %d) verification RMSE(u,v,h) : %f, %f, %f " % ( u_obs_number, v_obs_number, h_obs_number, u_RMSE,v_RMSE, h_RMSE)